package com.tyron.completion.java.compiler;

import javax.tools.Diagnostic;
import javax.tools.DiagnosticListener;
import javax.tools.JavaFileManager;
import javax.tools.JavaFileObject;
import com.sun.source.util.JavacTask;
import com.sun.source.util.TaskEvent;
import com.sun.source.util.TaskListener;
import com.sun.tools.javac.api.JavacTaskImpl;
import com.sun.tools.javac.api.JavacTool;
import com.sun.tools.javac.api.JavacTrees;
import com.sun.tools.javac.api.MultiTaskListener;
import com.sun.tools.javac.code.Types;
import com.sun.tools.javac.comp.Annotate;
import com.sun.tools.javac.comp.Check;
import com.sun.tools.javac.comp.CompileStates;
import com.sun.tools.javac.comp.Enter;
import com.sun.tools.javac.comp.Modules;
import com.sun.tools.javac.main.Arguments;
import com.sun.tools.javac.main.JavaCompiler;
import com.sun.tools.javac.model.JavacElements;
import com.sun.tools.javac.util.Context;
import com.sun.tools.javac.util.DefinedBy;
import com.sun.tools.javac.util.Log;
import com.tyron.common.logging.IdeLog;
import com.tyron.completion.java.compiler.services.CancelService;
import com.tyron.completion.java.compiler.services.NBAttr;
import com.tyron.completion.java.compiler.services.NBClassFinder;
import com.tyron.completion.java.compiler.services.NBEnter;
import com.tyron.completion.java.compiler.services.NBJavaCompiler;
import com.tyron.completion.java.compiler.services.NBJavacTrees;
import com.tyron.completion.java.compiler.services.NBLog;
import com.tyron.completion.java.compiler.services.NBMemberEnter;
import com.tyron.completion.java.compiler.services.NBParserFactory;
import com.tyron.completion.java.compiler.services.NBResolve;
import com.tyron.completion.java.compiler.services.NBTreeMaker;

import java.io.PrintWriter;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.logging.Handler;
import java.util.logging.LogRecord;
import java.util.logging.Logger;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

import me.xdrop.diffutils.DiffUtils;

/**
 * A pool of reusable JavacTasks. When a task is no valid anymore, it is returned to the pool, and its Context may be
 * reused for future processing in some cases. The reuse is achieved by replacing some components (most notably
 * JavaCompiler and Log) with reusable counterparts, and by cleaning up leftovers from previous compilation.
 *
 * <p>For each combination of options, a separate task/context is created and kept, as most option values are cached
 * inside components themselves.
 *
 * <p>When the compilation redefines sensitive classes (e.g. classes in the the java.* packages), the task/context is
 * not reused.
 *
 * <p>When the task is reused, then packages that were already listed won't be listed again.
 *
 * <p>Care must be taken to only return tasks that won't be used by the original caller.
 *
 * <p>Care must also be taken when custom components are installed, as those are not cleaned when the task/context is
 * reused, and subsequent getTask may return a task based on a context with these custom components.
 *
 * <p><b>This is NOT part of any supported API. If you write code that depends on this, you do so at your own risk. This
 * code and its internal interfaces are subject to change or deletion without notice.</b>
 */

public class ReusableCompiler {

    private static final Logger LOG = IdeLog.getCurrentLogger(ReusableCompiler.class);
    private static final JavacTool systemProvider = JavacTool.create();

    private final List<String> currentOptions = new ArrayList<>();
    private ReusableContext currentContext;
    private volatile boolean checkedOut;

    private final CancelServiceImpl cancelService = new CancelServiceImpl();

    public static class CancelServiceImpl extends CancelService {

        private final AtomicBoolean canceled = new AtomicBoolean(false);
        private final AtomicBoolean running = new AtomicBoolean(false);

        public void cancel() {
            canceled.set(true);
        }

        public boolean isRunning() {
            return running.get();
        }

        public void setRunning(boolean value) {
            running.set(value);
        }

        @Override
        public boolean isCanceled() {
            return false;
        }

        @Override
        protected void onCancel() {
            LOG.info("Compilation task cancelled.x");
            running.set(false);
        }
    }

    public CancelServiceImpl getCancelService() {
        return cancelService;
    }

    /**
     * Creates a new task as if by JavaCompiler and runs the provided worker with it. The
     * task is only valid while the worker is running. The internal structures may be reused from some previous
     * compilation.
     *
     * @param fileManager a file manager; if {@code null} use the compiler's standard filemanager
     * @param diagnosticListener a diagnostic listener; if {@code null} use the compiler's default method for reporting
     *     diagnostics
     * @param options compiler options, {@code null} means no options
     * @param classes names of classes to be processed by annotation processing, {@code null} means no class names
     * @param compilationUnits the compilation units to compile, {@code null} means no compilation units
     * @return an object representing the compilation
     * @throws RuntimeException if an unrecoverable error occurred in a user supplied component. The {@linkplain
     *     Throwable#getCause() cause} will be the error in user code.
     * @throws IllegalArgumentException if any of the options are invalid, or if any of the given compilation units are
     *     of other kind than {@linkplain JavaFileObject.Kind#SOURCE source}
     */
    public Borrow getTask(
		JavaFileManager fileManager,
		DiagnosticListener<? super JavaFileObject> diagnosticListener,
		Iterable<String> options,
		Iterable<String> classes,
		Iterable<? extends JavaFileObject> compilationUnits) {
        if (checkedOut) {
            throw new RuntimeException("Compiler is already in-use!");
        }
        checkedOut = true;
        List<String> opts =
			StreamSupport.stream(options.spliterator(), false)
                    .collect(Collectors.toList());
        if (!opts.equals(currentOptions)) {
            List<String> difference = new ArrayList<>(currentOptions);
            difference.removeAll(opts);
            LOG.warning("Options changed, creating new compiler \n difference: " + difference);
            currentOptions.clear();
            currentOptions.addAll(opts);
            currentContext = new ReusableContext(new ArrayList<>(opts), cancelService);
        }
        JavacTaskImpl task =
			(JavacTaskImpl)
			systemProvider.getTask(
			null, fileManager, diagnosticListener, opts, classes, compilationUnits, currentContext);
        task.addTaskListener(currentContext);

        cancelService.setRunning(true);

        return new Borrow(task, currentContext);
    }

    public ReusableContext getCurrentContext() {
        return currentContext;
    }

    public class Borrow implements AutoCloseable {
        final JavacTask task;
        boolean closed;

        Borrow(JavacTask task, ReusableContext ctx) {
            this.task = task;
        }

        @Override
        public void close() {
            if (closed) return;
            // not returning the context to the pool if task crashes with an exception
            // the task/context may be in a broken state
            currentContext.clear();
            try {
                Method method = JavacTaskImpl.class.getDeclaredMethod("cleanup");
                method.setAccessible(true);
                method.invoke(task);
            } catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
                throw new RuntimeException(e);
            } finally {
                checkedOut = false;
                closed = true;
            }
        }
    }

    static class ReusableContext extends Context implements TaskListener {

        List<String> arguments;

        ReusableContext(List<String> arguments, CancelService cancelService) {
            super();
            this.arguments = arguments;
            put(Log.logKey, ReusableLog.factory);
            put(JavaCompiler.compilerKey, ReusableJavaCompiler.factory);
            registerServices(this, cancelService);
        }

        private static void registerServices(Context context, CancelService cancelService) {
            NBAttr.preRegister(context);
            NBParserFactory.preRegister(context);
            NBTreeMaker.preRegister(context);
            NBJavacTrees.preRegister(context);
            NBResolve.preRegister(context);
            NBEnter.preRegister(context);
            NBMemberEnter.preRegister(context, false);
            NBClassFinder.preRegister(context);

            context.put(CancelService.cancelServiceKey, cancelService);
        }

        public void clear() {
            drop(Arguments.argsKey);
            drop(DiagnosticListener.class);
            drop(Log.outKey);
            drop(Log.errKey);
            drop(JavaFileManager.class);
            drop(JavacTask.class);
            drop(JavacTrees.class);
            drop(JavacElements.class);

            if (ht.get(Log.logKey) instanceof ReusableLog) {
                // log already inited - not first round
                ((ReusableLog) ReusableLog.instance(this)).clear();
                Enter.instance(this).newRound();
                ((ReusableJavaCompiler) ReusableJavaCompiler.instance(this)).clear();
                Types.instance(this).newRound();
                Check.instance(this).newRound();
                Modules.instance(this).newRound();
                Annotate.instance(this).newRound();
                CompileStates.instance(this).clear();
                MultiTaskListener.instance(this).clear();
            }
        }

        @Override
        @DefinedBy(DefinedBy.Api.COMPILER_TREE)
        public void finished(TaskEvent e) {
            // do nothing
        }

        @Override
        @DefinedBy(DefinedBy.Api.COMPILER_TREE)
        public void started(TaskEvent e) {
            // do nothing
        }

        <T> void drop(Key<T> k) {
            ht.remove(k);
        }

        <T> void drop(Class<T> c) {
            ht.remove(key(c));
        }

        /**
         * Reusable JavaCompiler; exposes a method to clean up the component from leftovers associated with previous
         * compilations.
         */
        static class ReusableJavaCompiler extends NBJavaCompiler {

            static final Factory<JavaCompiler> factory = ReusableJavaCompiler::new;

            ReusableJavaCompiler(Context context) {
                super(context);
            }

            @Override
            public void close() {
                // do nothing
            }

            void clear() {
                newRound();
            }

            @Override
            protected void checkReusable() {
                // do nothing - it's ok to reuse the compiler
            }
        }

        /**
         * Reusable Log; exposes a method to clean up the component from leftovers associated with previous
         * compilations.
         */
        static class ReusableLog extends Log {

            private static final Factory<Log> factory = ReusableLog::new;

            Context context;

            ReusableLog(Context context) {
                super(context);
                this.context = context;
            }

            void clear() {
                recorded.clear();
                sourceMap.clear();
                nerrors = 0;
                nwarnings = 0;
                // Set a fake listener that will lazily lookup the context for the 'real' listener. Since
                // this field is never updated when a new task is created, we cannot simply reset the field
                // or keep old value. This is a hack to workaround the limitations in the current infrastructure.
                diagListener =
					new DiagnosticListener<JavaFileObject>() {
					DiagnosticListener<JavaFileObject> cachedListener;

					@Override
					@DefinedBy(DefinedBy.Api.COMPILER)
					@SuppressWarnings("unchecked")
					public void report(Diagnostic<? extends JavaFileObject> diagnostic) {
						if (cachedListener == null) {
							cachedListener = context.get(DiagnosticListener.class);
						}
						cachedListener.report(diagnostic);
					}
				};
            }
        }
    }
}
